# -*- coding: utf-8 -*-
import os
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import PySimpleGUI as sg
import matplotlib
import mat_util

matplotlib.use('TkAgg')

def delete_figure_agg(figure_agg):
    figure_agg.get_tk_widget().forget()
    plt.close('all')

def draw_figure(canvas, figure):
    figure_canvas_agg = FigureCanvasTkAgg(figure, canvas)
    figure_canvas_agg.draw()
    figure_canvas_agg.get_tk_widget().pack(side='top', fill='both', expand=1)
    return figure_canvas_agg

def plot_single_1dimg(img_data, img_shape=(64, 64), imshow_origin='lower'):
    '''
    单张图像为img_w*img_h的1D展开图像
    :param img_data: 1d array, 图像数据，形如：(img_w * img_h, )
    :param img_shape: tuple, 显示的图像高和宽，形如：(img_h, img_w)
    :param imshow_origin: 显示图像时的原点位置，取'lower'或'upper'
    :return:
    '''
    img_show = img_data.reshape(img_shape)
    plt.subplot(1, 1, 1)
    plt.imshow(img_show, 'gray', origin=imshow_origin)
    return plt.gcf()

def plot_single_2dimg(img_data, imshow_origin='lower'):
    '''
    单张图像为(img_w,img_h)的2D图像
    :param img_data: 2d array, 图像数据，形如：(img_w, img_h)
    :param imshow_origin: 显示图像时的原点位置，取'lower'或'upper'
    :return:
    '''
    plt.subplot(1, 1, 1)
    plt.imshow(img_data, 'gray', origin=imshow_origin)
    return plt.gcf()

if __name__ == '__main__':
    # define the window layout
    col2 = sg.Column([
        [sg.Frame('修正当前数据', [
            [
            sg.Radio('shift在变换之前', group_id='-group_shift_first-'),
            sg.Radio('shift在变换之后', group_id='-group_shift_first-'),
            sg.Button('修正当前', key='-btn_revise-')
            ]
        ])],
    ])
    layout = [[sg.Text('第0幅图像', key='-text_imageno-')],
              [sg.Canvas(key='-CANVAS-')],
              [sg.Text('图像大小', key='-text_image_width-'),
               sg.Combo(values=[64,96,128,256,384,512],
                        default_value=64,
                        enable_events=True,
                        size=(20, 1), k='-comb_img_width-')],
              [sg.Text('图像原点', key='-text_imageorigin-'),
               sg.Combo(values=['左下角', '左上角'],
                       default_value='左下角',
                        enable_events=True,
                       size=(20, 1), k='-comb_origin-')],
              [sg.Text('图像形式', key='-text_imageforms-'),
               sg.Combo(values=['1D展开并连接的频域数据', '2D频域数据', '1D图像数据', '2D图像数据'],
                        default_value='2D图像数据',
                        enable_events=True,
                        size=(20, 1), k='-comb_imageforms-')],
              [sg.Button('打开目标图像mat', key='-btn_open-'),
               sg.Button('下一个', key='-btn_next-'), sg.Button('上一个', key='-btn_previous-')],
              [sg.Text('跳转编号（0开始）：'), sg.InputText('0', key='-in_goto_imgnumber-', size=(5,1)),
               sg.Button('跳转',key='-btn_goto_img-')],
              [sg.Frame('修正当前数据', [
                  [
                      sg.Text('shift零频率：'),
                      sg.Radio('是', key='-radio_shift_yes-', group_id='-group_whethershift-', default=True ),
                      sg.Radio('否', key='-radio_shift_no-', group_id='-group_whethershift-')
                  ],
                  [
                      sg.Text('shift位置：'),
                      sg.Radio('变换之前', key='-radio_shiftfirst-', group_id='-group_shift_first-',default=True,),
                      sg.Radio('变换之后', key='-radio_shiftlast-', group_id='-group_shift_first-')
                  ],
                  [
                      sg.Button('修正当前', key='-btn_preview_revise-'), sg.Button('应用修正', key='-btn_apply_revise-')
                  ]
              ])],
              ]
    # layout=[
    #     [col2]
    # ]
    # create the form and show it without the plot
    window = sg.Window('Mat图像查看器', layout, finalize=True, element_justification='center', font='宋体 14')

    mat_data_revised = {}   # 修改后的数据放到这里，类似于{img_no: arr_data, img_no2: arr_data2}
    fig_canvas_agg = None  # 绘图的对象
    current_img_number = 0  #当前显示图像的编号,编号从0开始
    mat_image_data = None #从mat中加载的图像数据
    total_nb_imgs = 0
    imshow_origin = 'lower'
    img_shape = (64,64)
    while True:
        event, values = window.read()
        if fig_canvas_agg:
            # ** IMPORTANT ** Clean up previous drawing before drawing again
            delete_figure_agg(fig_canvas_agg)
        if event == sg.WINDOW_CLOSED:
            break
        elif event=='-btn_open-':
            current_img_number = 0
            mat_image_data = None
            total_nb_imgs = 0
            loaded_mat_file = sg.popup_get_file('请选择mat文件',
                                                '请选择mat文件',
                                                '',
                                                '.mat',
                                                file_types=(("Mat Files", ". mat"),)
                                                )
            if (loaded_mat_file is not None) and os.path.isfile(loaded_mat_file):
                mat_image_data = load_mat(loaded_mat_file)
                print(f'{mat_image_data.shape=}')
                try:
                    if values['-comb_imageforms-'] == '1D图像数据':
                        _, total_nb_imgs = mat_image_data.shape  #获取总的图像数目
                    elif values['-comb_imageforms-']=='2D图像数据':
                        img_h, img_w, total_nb_imgs = mat_image_data.shape
                except:
                    window['-text_imageno-'].update('打开图像出错！请检查图像格式是否选择正确')
                    continue
        elif event=='-btn_next-':
            # 绘制下一幅图像
            current_img_number = current_img_number + 1
            if current_img_number>total_nb_imgs-1:
                current_img_number=total_nb_imgs-1
        elif event=='-btn_previous-':
            # 绘制前一图像
            current_img_number = current_img_number - 1
            if current_img_number<0:
                current_img_number=0
        elif event=='-btn_goto_img-':
            # 跳转到指定图像
            imgnumber_txt = values['-in_goto_imgnumber-']
            if imgnumber_txt.isnumeric():
                current_img_number = int(imgnumber_txt)
                if current_img_number<0:
                    current_img_number=0
                elif current_img_number>total_nb_imgs-1:
                    current_img_number=total_nb_imgs-1
        elif event=='-comb_origin-':
            if values['-comb_origin-']=='左上角':
                imshow_origin='upper'
            elif values['-comb_origin-']=='左下角':
                imshow_origin = 'lower'
        elif event == '-comb_img_width-':
            img_shape = (values['-comb_img_width-'], values['-comb_img_width-'])
        elif event=='-btn_preview_revise-':
            # 预览修正当前图像
            cur_img2d = mat_image_data[..., current_img_number]
            # 图像转到频域，转到频域时需要移频，因此这里whether_shift一直设置为True
            cur_freq2d = mat_util.get_2dfreqdata_from2dImage(cur_img2d, whether_shift=values['-radio_shift_yes-'])
            # 频域转回图像
            back_img2d = mat_util.get_2dImage_from2dFreq(cur_freq2d, whether_shift=values['-radio_shift_yes-'],
                                                  shift_first=values['-radio_shiftfirst-'])
            # 添加到已经修正过的列表中
            mat_data_revised[current_img_number] = np.abs(back_img2d)
        elif event=='-btn_apply_revise-':
            pass
        # 展示第current_img_number幅图像
        if mat_image_data is not None:
            # 如修正列表中有当前编号的图像，则【优先显示修正列表中的图像】
            cur_data = mat_data_revised.get(current_img_number)
            # 如果修正列表中没有当前编号图像，则从原始数据中获取当前编号的图像数据
            if cur_data is None:
                cur_data = mat_image_data[..., current_img_number]
            if values['-comb_imageforms-']=='1D图像数据':
                first_fig = plot_single_1dimg(cur_data, img_shape, imshow_origin)
            elif values['-comb_imageforms-']=='2D图像数据':
                first_fig = plot_single_2dimg(cur_data, imshow_origin)
            fig_canvas_agg = draw_figure(window['-CANVAS-'].TKCanvas, first_fig)
            window['-text_imageno-'].update(f'第{current_img_number}/{total_nb_imgs-1}幅图像')
    window.close()